library(R2jags)
library(xlsx)
library(openxlsx)
library(spatstat)
library(dplyr)
library(readr)
require(R2wd)
library(LaplacesDemon) # for Geweke.Diagnostic()
library(ggplot2)

options(max.print=999999999)
memory.limit(10000);memory.size(max=T)

getparams<-function(jagssum,parameters=c('mu','sigma','diff','disc','beta'),questions=questions){
  for (n in parameters){v<-data.frame(jagssum[grepl(n,dimnames(jagssum)[[1]]),"mean"])
    s<-data.frame(jagssum[grepl(n,dimnames(jagssum)[[1]]),"sd"])
    for (w in c(parameters,'\\[','\\]')){row.names(v)<-gsub(w,"",row.names(v));row.names(s)<-gsub(w,"",row.names(s))}
    v <- cbind(rownames(v),v);s<-cbind(rownames(s),s)
    if (n=='beta'){v<-data.frame(v[,1:2]);s<-data.frame(s[,1:2]);colnames(v)<-c('n','mean.beta');colnames(s)<-c('n','sd.beta');pq<-left_join(v,s) 
    } else {if (n=="mu"|n=="sigma"){colnames(v)<-c('t',paste0("mean.",n));colnames(s)<-c('t',paste0("sd.",n));pq<-left_join(v,s)
      } else {if (n=="diff"){colnames(v)<-c('qn',paste0('mean.',n));colnames(s)<-c('qn',paste0('sd.',n))
          if (nrow(v)==info$nquest){v$qn<-as.character(v$qn);s$qn<-as.character(s$qn)
          } else {v$b<-substr(v$qn,(nchar(as.character(v$qn))+1)-1,nchar(as.character(v$qn)))
            s$b<-substr(s$qn,(nchar(as.character(s$qn))+1)-1,nchar(as.character(s$qn)))
            row.names(v)<-c(1:nrow(v));row.names(s)<-c(1:nrow(s))
            v$qn<-gsub(",1|,2|,3","",v$qn)
            s$qn<-gsub(",1|,2|,3","",s$qn)}} else {colnames(v)<-c('qn',paste0('mean.',n));colnames(s)<-c('qn',paste0('sd.',n))
          v$qn<-as.character(v$qn);s$qn<-as.character(s$qn)}
        pq<-left_join(v,s);pq<-left_join(pq,questions);pq$question<-as.character(pq$question)}}
    return(pq)}}

setwd(dirname(rstudioapi::getActiveDocumentContext()$path)) # set wd to source file location

############################################################################################

## JAGS code here
main_model="data{for (i in 1:len){opinionr[i,2]<-round(support[i]*samplesize[i]/100) 
		opinionr[i,1]<-round(neutral[i]*samplesize[i]/100)}}
model{for (i in 1:len){for(j in 1:2){opinionr[i,j]~dbin(p[i,j],samplesize[i])      
		p[i,j]~dbeta(alpha[i,j], beta)              
		alpha[i,j]<--beta*m[i,j]/(m[i,j]-1)         
		m[i,j]<-phi(x[i,j])             
	  x[i,j]<-(mu[t[i]]-diff[qn[i],j])/sqrt((disc[qn[i]])^2+(sigma[t[i]])^2)}}
beta~dunif(0,100)
mu[begin]~dnorm(0,1)
sigma_mu~dunif(0,10)
for (i in 1:nquest){disc[i]~dunif(0,6) 
	for (j in 1:2){diff0[i,j]~dunif(-4,4)}
	diff[i,1:2]<-sort(diff0[i,])}
for (i in begin:end){sigma[i]~dunif(0,10)}
for (i in begin:(end-1)){mu_raw[i]~dnorm(0,1)}
for (i in (begin+1):end){mu[i]~dnorm(mu[i-1]+sigma_mu*mu_raw[i-1],100)}}"
## Save to temp file
fileConn=file("main_model.temp")
writeLines(main_model,fileConn)
close(fileConn)

###########################################################################################

outdir<-'irt_estimates/';dir.create(file.path(getwd(),outdir),showWarnings=F)
inpdir<-'raw_indicators/'

mainsheet<-read.table(paste0(inpdir,'EU_irt.txt'),sep=';',header=T)  # Read in data
mainsheet<-subset(mainsheet,as.numeric(gsub(' S1| S2','',mainsheet$t))>=1973)
mainsheet$t<-as.numeric(mainsheet$t)
mainsheet$asked<-NULL
mainsheet<-suppressMessages(left_join(mainsheet,mainsheet %>% group_by(question) %>% tally() %>% as.data.frame()))
names(mainsheet)[names(mainsheet)=='n']<-'asked'  
mainsheet<-subset(mainsheet,mainsheet$asked>1) 

policy<-list('jags','social','econom','market','asylum') # Delete or keep accordingly

for (p in policy){sheet<-mainsheet;sheet$policy<-sheet$thetopic;print(p)
  wb<-paste0(outdir,'EU_irt_estimates.xlsx');if (p!='jags'){sheet<-sheet[sheet$policy==p,]}
  
  rownames(sheet) <- 1:nrow(sheet);sheet$question<-factor(sheet$question) # Reset row numbers and update levels of q names after subsetting data
  sheet$qn<-as.numeric(sheet$question);sheet<-sheet[order(sheet$qn,sheet$t),] # Sort by q number and time
  questions<-unique(sheet[c("qn","question")]);questions$question<-as.character(questions$question);questions$qn<-as.character(questions$qn)
  info<-list(len=nrow(sheet),nquest=length(unique(sheet$qn)),begin=min(sheet$t),end=max(sheet$t));data <-c(sheet,info)
  
  start_time<-Sys.time();print(start_time);print(paste0('Using ',info$nquest,' questions and ',info$len,' observations.'))
  jagsfit<-jags.parallel(data=data,parameters.to.save=c('mu','sigma',"diff","disc","beta"),
                         n.burnin=10000,n.iter=100000,n.thin=1,n.chains=3,
                         model.file="main_model.temp",jags.seed=123)
  duration<-Sys.time()-start_time;print(duration)
  
  fit.mcmc<-as.mcmc(jagsfit);print(gelman.diag(fit.mcmc,autoburnin=F,transform=F)$mpsrf)
  
  sink(file=paste0(outdir,'model_summaries/summary_model_EU_',p,'.txt'));print(duration);print(jagsfit);sink()
  sink(file=paste0(outdir,'model_summaries/gelman_diagnostic_EU_',p,'.txt'));print(gelman.diag(fit.mcmc,autoburnin=F));sink() # Gelman diagnostic: prsf should be below 1.1
  sink(file=paste0(outdir,'model_summaries/geweke_diagnostic_EU_',p,'.txt'));print(Geweke.Diagnostic(fit.mcmc));sink() # Geweke diagnostic: should be between -2 and 2
  
  for (par in c('mu','sigma','diff','disc','beta')){p.func<-getparams(jagsfit$BUGSoutput$summary,par,questions); assign(par,p.func)};print(" ")
  
  write.table(beta,file=paste0(outdir,'model_parameters/beta_EU_',p,'.txt'),row.names=F,quote=F,sep="\t")
  write.table(diff,file=paste0(outdir,'model_parameters/difficulty_EU_',p,'.txt'),row.names=F,quote=F,sep="\t")
  write.table(disc,file=paste0(outdir,'model_parameters/discrimination_EU_',p,'.txt'),row.names=F,quote=F,sep="\t")
  
  df<-merge(mu,sigma);df$t<-as.numeric(as.character(df$t));df<-df[order(df$t),] 
  if ((nrow(diff))==info$nquest){diff<-diff} else {j=as.character(max(diff$b));diff<-subset(diff,b==j)}
  for (r in 1:nrow(df)){
    df$Opinion[r]<-pnorm((df$mean.mu[r]-mean(diff$mean.diff))/(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5,0,1)*100
    df$Opinion_minus2sd[r]<-pnorm(((df$mean.mu[r]-2*df$sd.mu[r])-mean(diff$mean.diff))/(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5,0,1)*100
    df$Opinion_plus2sd[r]<-pnorm(((df$mean.mu[r]+2*df$sd.mu[r])-mean(diff$mean.diff))/(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5,0,1)*100
    df$Opinion_minussigma[r]<-df$Opinion[r]-(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5
    df$Opinion_plussigma[r]<-df$Opinion[r]+(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5}
  if (!file.exists(wb)){workbook<-createWorkbook();addWorksheet(workbook,p);writeData(workbook,sheet=p,x=df,rowNames=F)
    saveWorkbook(workbook,wb,overwrite=T)} else {workbook<-loadWorkbook(wb)
    if (p %in% names(workbook)){removeWorksheet(workbook,p);addWorksheet(workbook,p)} else {addWorksheet(workbook,p)}
    writeData(workbook,sheet=p,x=df,rowNames=F);saveWorkbook(workbook,wb,overwrite=T)}}
